#!/usr/bin/python3
# -*- coding: utf-8 -*-

import cv2
import numpy as np

np.set_printoptions(formatter={'float': '{: 0.4f}'.format})

import rospy
import message_filters

from geometry_msgs.msg import Transform,Pose,Point32

from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError


from bs_imu import bs_imu
from bs_imu import bs_gps_raw
from bs_img import bs_img_real
from bs_img import bs_img_base
from bs_img.bs_cfg_real_5m import *
from bs_img.bs_lie_solve import *
from bs_img.serial_pack import *

import time
'''
输入图像
输出相对的解算值
'''




class BS_Solver:
    def __init__(self):
        lf_img_topic = "/camera_lf/image_dh"
        self.last_plane_pts = np.array([])   #保存上一帧点用于配对
        self.is_first_frame = True                  #判断是否为第一帧

        self.img_bridge = CvBridge()

        self.invalid_frame_num_plane = 0      #compute continus invalid frame
        self.invalid_frame_num_drogue = 0
        self.invalid_frame_max = 8     #if invaild frame < invalid_frame_max, output results of the last frame

        #Kalman Filter
        self.Q = np.eye(3)
        self.R = np.eye(3)
        self.x_plane = np.array([np.nan, np.nan, np.nan])
        self.w_plane = 1.5
        self.x_drogue = np.array([np.nan, np.nan, np.nan])
        self.w_drogue = 1.5

        self.lf_img_sub = rospy.Subscriber(lf_img_topic, Image, self.lf_img_cb)

        self.lf_plane_pub = rospy.Publisher("/bs_debuger/result/lf_plane", Transform , queue_size=3)
        self.lf_drogue_pub = rospy.Publisher("/bs_debuger/result/lf_drogue", Transform, queue_size=3)





    def lf_img_cb(self,img_msg):
        img_gray = self.img_bridge.imgmsg_to_cv2(img_msg, img_msg.encoding)
  
        _, img_bin = cv2.threshold(img_gray, 70, 0xff, cv2.THRESH_BINARY)
        # 如果相机倒置需要旋转一下
        # lf_img_gray = cv2.rotate(lf_img_gray, cv2.ROTATE_180)
        plane_t_vec = [np.nan]*3
        drogue_t_vec = [np.nan]*3
        drogue_pts = []
        
        gt_pts,aera_list = bs_img_real.img_to_pts(img_bin)
        gt_pts = np.array(gt_pts)
        cnt_plane_pts = [ ]


        if self.is_first_frame:   #第一帧图像
            drogue_result_dict = bs_img_real.solve_drogue(gt_pts, camK_lf, taper_radius=r_drogue)        
            if drogue_result_dict["img_valid"]:
                drogue_t_vec = drogue_result_dict["t_vec"]
                drogue_pts = drogue_result_dict["drogue_px_calcL"]
            else:
                drogue_pts = []
            drogue_pts = np.array(drogue_pts)
            if np.size(drogue_pts) >0:
                ind = ~np.isin(gt_pts[:, 0], drogue_pts[:, 0]) | ~np.isin(gt_pts[:, 1], drogue_pts[:, 1])
                plane_pts = np.array(gt_pts)[ind]
                ind2 = ind == False
                plane_aeras = np.array(aera_list)[ind2] 
            else:
                plane_pts = gt_pts
                plane_aeras = np.array(aera_list)
            if len(plane_pts)>=5:
                plane_result_dict = bs_img_base.solve_plane_pt_new(plane_pts,plane_real_ptL,cam_K=camK_lf,distort_coeffs=distortion_coeffs_lf,plane_area_ptL=plane_aeras)
                if plane_result_dict["img_valid"]:
                    self.is_first_frame = False 
                    cnt_plane_pts = plane_result_dict["plane_px_calcL"]
                    plane_t_vec = plane_result_dict["t_vec"].flatten()
                    plane_quat = bs_imu.rot_vec_to_quat(plane_result_dict["r_vec"].flatten())

        else:   #除第一帧以外的图像
            match_pt_idxs = find_matched_pts(self.last_plane_pts,gt_pts)
            flag_empty = True   #指示匹配是否为空
            plane_aeras = []
            for match_pt_idx in match_pt_idxs:
                if match_pt_idx is not None:
                    cnt_plane_pts.append(gt_pts[match_pt_idx])
                    plane_aeras.append(gt_pts[match_pt_idx])
                    flag_empty = False
            if flag_empty:
                self.is_first_frame = True   #匹配为空重新开始，视为第一帧
            
            # 至少有5个点则解算飞机特征点
            if len(cnt_plane_pts) >= 5:    
                plane_result_dict = bs_img_base.solve_plane_pt_new(cnt_plane_pts,plane_real_ptL,cam_K=camK_lf,distort_coeffs=distortion_coeffs_lf,plane_area_ptL=plane_aeras)
                if plane_result_dict["img_valid"]:
                    self.is_first_frame = False 
                    cnt_plane_pts = plane_result_dict["plane_px_calcL"]
                    plane_t_vec = plane_result_dict["t_vec"].flatten()
                    # plane_quat = bs_imu.rot_vec_to_quat(plane_result_dict["r_vec"].flatten())
                else:
                    cnt_plane_pts = []
            else:
                cnt_plane_pts = []
            cnt_plane_pts = np.array(cnt_plane_pts)
                
            if np.size(cnt_plane_pts) >0:
                ind = ~np.isin(gt_pts[:, 0], cnt_plane_pts[:, 0]) | ~np.isin(gt_pts[:, 1], cnt_plane_pts[:, 1])
                drogue_pts = np.array(gt_pts)[ind]
            else:
                drogue_pts = gt_pts
            
            drogue_result_dict = bs_img_real.solve_drogue(drogue_pts, camK_lf, taper_radius=r_drogue)        
            if drogue_result_dict["img_valid"]:
                drogue_t_vec = drogue_result_dict["t_vec"]

        
        self.last_plane_pts = cnt_plane_pts - np.mean(cnt_plane_pts,axis=0)


        plane_is_visiable = True if not(np.all(np.isnan(plane_t_vec))) else False
        drogue_is_visiable = True if not(np.all(np.isnan(drogue_t_vec))) else False
        
        #compensate
        if plane_is_visiable:
            plane_t_vec[0] = plane_t_vec[0]*0.9336+0.23
            plane_t_vec[1] = plane_t_vec[1]+0.32
            plane_t_vec[2] = plane_t_vec[2]*0.9396+3.3
        if drogue_is_visiable:
            drogue_t_vec[0] = drogue_t_vec[0]*0.9336+0.23
            drogue_t_vec[1] = drogue_t_vec[1]+0.32
            drogue_t_vec[2] = drogue_t_vec[2]*0.9396+3.3


        # KF for plane
        if not(plane_is_visiable):
            self.invalid_frame_num_plane = self.invalid_frame_num_plane + 1
            if self.invalid_frame_num_plane <= self.invalid_frame_max:
                plane_res = self.x_plane
                if not(np.all(np.isnan(plane_res))):
                    plane_res_frame = serial_frame("drogue", True, round(plane_res[2]*100),round(plane_res[0]*100),round(plane_res[1]*100))  
                else:
                    plane_res_frame = serial_frame("drogue",False,0,0,0)
            else:
                plane_res_frame = serial_frame("drogue",False,0,0,0)
        
        else:
            self.invalid_frame_num_plane = 0

            if np.all(np.isnan(self.x_plane)):   #the first frame
                self.x_plane = plane_t_vec
            else:
                miu = plane_t_vec-self.x_plane
                self.x_plane = self.x_plane + np.matmul(self.Q,np.matmul(np.linalg.pinv(self.Q+self.R/self.w_plane),miu.T))
                self.w_plane = 1.5/(1+np.matmul(miu,np.matmul(np.linalg.pinv(self.R),miu.T)))
            plane_res_frame = serial_frame("drogue",True,round(self.x_plane[2]*100),round(self.x_plane[0]*100),round(self.x_plane[1]*100))  

        uart.write(plane_res_frame)
        time.sleep(0.0001)

        # # KF for drogue
        # if not(drogue_is_visiable):
        #     self.invalid_frame_num_drogue = self.invalid_frame_num_drogue + 1
        #     if self.invalid_frame_num_drogue <= self.invalid_frame_max:
        #         drogue_res = self.x_drogue
        #         if not(np.all(np.isnan(drogue_res))):
        #             drogue_res_frame = serial_frame("plane", True, round(drogue_res[2]*100),round(drogue_res[0]*100),round(drogue_res[1]*100))  
        #         else:
        #             drogue_res_frame = serial_frame("plane",False,0,0,0)
        #     else:
        #         drogue_res_frame = serial_frame("plane",False,0,0,0)
        
        # else:
        #     self.invalid_frame_num_drogue = 0

        #     if np.all(np.isnan(self.x_drogue)):   #the first frame
        #         self.x_drogue = drogue_t_vec
        #     else:
        #         miu = drogue_t_vec-self.x_drogue
        #         self.x_drogue = self.x_drogue + np.matmul(self.Q,np.matmul(np.linalg.pinv(self.Q+self.R/self.w_drogue),miu.T))
        #         self.w_drogue = 1.5/(1+np.matmul(miu,np.matmul(np.linalg.pinv(self.R),miu.T)))
        #     drogue_res_frame = serial_frame("plane",True,round(self.x_drogue[2]*100),round(self.x_drogue[0]*100),round(self.x_drogue[1]*100))  

        # uart.write(drogue_res_frame)
        # time.sleep(0.0001)

        plane_msg = Transform()
        plane_msg.translation.x = plane_t_vec[0]
        plane_msg.translation.y = plane_t_vec[1]
        plane_msg.translation.z = plane_t_vec[2]
        # plane_msg.rotation.w = plane_quat[0]
        # plane_msg.rotation.x = plane_quat[1]
        # plane_msg.rotation.y = plane_quat[2]
        # plane_msg.rotation.z = plane_quat[3]
        self.lf_plane_pub.publish(plane_msg)
        
        drogue_msg = Transform()
        drogue_msg.translation.x = drogue_t_vec[0]
        drogue_msg.translation.y = drogue_t_vec[1]
        drogue_msg.translation.z = drogue_t_vec[2]
        self.lf_drogue_pub.publish(drogue_msg)


        print(plane_t_vec)



if __name__ == "__main__":
    rospy.init_node('bs_solver_node', anonymous=True)
    port, baud =  '/dev/ttyUSB0', 115200
    uart = serial.Serial(port,baud,parity=serial.PARITY_NONE)
    bs_solver = BS_Solver()
    rospy.spin()
